import json

import jqdatasdk
from datetime import datetime
from stocksys import models as stocksys_models
import pickle


class PredictionMode:
    data = None
    clf = None
    prediction_function = None
    feature_engineering = None
    output_function = None
    fit_function = None

    def __init__(self, model):
        predict_model = __import__(model)
        predict_model = predict_model.quant_rules
        rule_name = model.split(".")[-1]
        feature_engineering = predict_model.__getattribute__(rule_name).feature_engineering
        fit_function = predict_model.__getattribute__(rule_name).fit_function
        prediction_function = predict_model.__getattribute__(rule_name).prediction_function
        output_function = predict_model.__getattribute__(rule_name).output_function
        self.feature_engineering = feature_engineering
        self.fit_function = fit_function
        self.prediction_function = prediction_function
        self.output_function = output_function

    def init_base_data(self, fullcode, start_date: datetime, end_date: datetime, exchange_type="stock",
                       method="JQDATA"):

        data = None
        if method == "JQDATA":
            api_keys = stocksys_models.APIKeys.objects.all().first()
            jqdatasdk.auth(api_keys.username, api_keys.password)
            security = jqdatasdk.get_security_info(fullcode)
            count = abs(
                (datetime.strptime(start_date, "%Y-%m-%d") - datetime.strptime(end_date, "%Y-%m-%d")).days)
            # print('count:%s \n end_date:%s \n start_date:%s' % (count,end_date,start_date))
            if exchange_type == "fund":  # TODO:加载基金数据
                data = jqdatasdk.attribute_history_engine(end_dt=end_date, count=count, security=security, unit='1d',
                                                          fields=('open', 'close', 'high', 'low', 'volume', 'money'),
                                                          skip_paused=True,
                                                          df=True, fq=None)
            elif exchange_type == "stock":  # TODO:加载股票数据
                data = jqdatasdk.get_price(security, start_date=start_date, end_date=end_date,
                                           frequency='daily',
                                           fields=None,
                                           skip_paused=False, fq='pre')
            elif exchange_type == "market_index":  # TODO:加载大盘数据
                data = jqdatasdk.get_price(start_date=start_date,
                                           end_date=end_date,
                                           security=security,
                                           frequency='daily',
                                           skip_paused=True,
                                           fq='pre')
            elif exchange_type == "industry":  # TODO:加载行业数据
                data = jqdatasdk.attribute_history_engine(end_dt=end_date, count=count, security=security, unit='1d',
                                                          fields=(
                                                              'open', 'close', 'high', 'low', 'volume', 'money'),
                                                          skip_paused=True,
                                                          df=True, fq=None)
            elif exchange_type == "conceptual":  # TODO:加载概念数据
                data = jqdatasdk.attribute_history_engine(end_dt=end_date, count=count, security=security, unit='1d',
                                                          fields=(
                                                              'open', 'close', 'high', 'low', 'volume', 'money'),
                                                          skip_paused=True,
                                                          df=True, fq=None)

            data['date'] = data.index
            rows, cols = data.shape
            data.index = [x for x in range(rows)]
            self.data = data
            jqdatasdk.logout()
        self.data = data[['open', 'high', 'low', 'close', 'volume','date']]

    def default_prediction_function(self, context):
        clf = None
        accuracy = 0
        return clf, accuracy

    def default_fit_function(self, context):
        clf = None
        accuracy = 0
        return clf, accuracy

    def default_feature_engineering(self, context):
        return context

    def default_output_function(self, context):
        return context

    def predict(self):
        pass

    def save_model(self, scale):
        pass

    def train_save(self, scale, model_database: stocksys_models.ManagedCode):
        start_time = datetime.now()
        context = dict()
        context['args'] = dict()
        context['data'] = self.data
        model_database.update_datetime = start_time

        # TODO:调用特征工程函数
        context = self.feature_engineering(context)
        context['scale'] = scale
        # TODO:调用训练函数
        context = self.fit_function(context)

        model_database.train_time = (datetime.now() - start_time).seconds
        model_database.model_file = pickle.dumps(context['clf'])
        model_database.accuracy = context['accuracy']
        model_database.model_paramater = pickle.dumps(json.dumps(context['args']))
        model_database.save()

    def train_predict(self, scale):
        context = dict()
        context['args'] = dict()
        context['data'] = self.data
        # TODO:调用特征工程函数
        context = self.feature_engineering(context)
        # TODO:调用训练函数
        context['scale'] = scale
        context = self.fit_function(context)
        # TODO:调用预测函数
        context = self.prediction_function(context)
        return self.output_function(context)

    def train_predict_save(self, scale, model_database: stocksys_models.ManagedCode):
        start_time = datetime.now()
        context = dict()
        context['args'] = dict()
        context['data'] = self.data
        # print(self.data)
        model_database.update_datetime = start_time

        # TODO:调用特征工程函数
        context = self.feature_engineering(context)
        # TODO:调用训练函数
        context['scale'] = scale
        context = self.fit_function(context)
        # TODO:调用预测函数
        context = self.prediction_function(context)

        model_database.train_time = (datetime.now() - start_time).seconds
        model_database.model_file = pickle.dumps(context['clf'])
        model_database.accuracy = context['accuracy']
        model_database.model_paramater = json.dumps(context['args'])
        model_database.save()
        return self.output_function(context)
